import os
import cv2
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.neighbors import KNeighborsClassifier
from sklearn.svm import SVC
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import accuracy_score
from sklearn.metrics import classification_report, accuracy_score, precision_score, recall_score, f1_score
import pandas as pd

# HOG 特征提取器
def get_hog_descriptor():
    return cv2.HOGDescriptor(
        _winSize=(64, 128), # HOG 输入大小
        _blockSize=(16, 16),# 块大小
        _blockStride=(8, 8),# 块步长
        _cellSize=(8, 8),   # 细胞大小
        _nbins=12      # 直方图通道数
    )

# 加载数据集
def load_inria_dataset(pos_dir, neg_dir):
    data = []
    labels = []
    hog = get_hog_descriptor()

    # 加载正样本（行人）
    for filename in os.listdir(pos_dir):
        filepath = os.path.join(pos_dir, filename)
        img = cv2.imread(filepath, cv2.IMREAD_GRAYSCALE)
        if img is not None:
            img = cv2.resize(img, (64, 128))  # 调整为 HOG 输入大小
            features = hog.compute(img).flatten()
            data.append(features)
            labels.append(1)  # 行人标签为 1

    # 加载负样本（非行人）
    for filename in os.listdir(neg_dir):
        filepath = os.path.join(neg_dir, filename)
        img = cv2.imread(filepath, cv2.IMREAD_GRAYSCALE)
        if img is not None:
            img = cv2.resize(img, (64, 128))  # 调整为 HOG 输入大小
            features = hog.compute(img).flatten()
            data.append(features)
            labels.append(0)  # 非行人标签为 0

    return np.array(data), np.array(labels)

# 主函数
def main():
    # 数据集路径
    pos_dir = "./INRIAPerson/Train/pos"  # 正样本路径
    neg_dir = "./INRIAPerson/Train/neg"  # 负样本路径

    # 加载数据
    print("加载数据集...")
    data, labels = load_inria_dataset(pos_dir, neg_dir)

    # 标准化数据
    print("标准化数据...")
    scaler = StandardScaler()
    data = scaler.fit_transform(data)

    # 划分训练集和测试集
    print("划分训练集和测试集...")
    X_train, X_test, y_train, y_test = \
        train_test_split(data, labels, test_size=0.2, random_state=42)

    # 使用 KNN 进行分类
    print("训练 KNN 模型...")
    knn = KNeighborsClassifier(n_neighbors=5, weights='distance')
    knn.fit(X_train, y_train)

    # 测试 KNN 模型
    print("测试 KNN 模型...")
    y_pred_knn = knn.predict(X_test)
    accuracy_knn = accuracy_score(y_test, y_pred_knn)
    precision_knn = precision_score(y_test, y_pred_knn)
    recall_knn = recall_score(y_test, y_pred_knn)
    f1_knn = f1_score(y_test, y_pred_knn)

    # 使用 SVM 进行分类
    print("训练 SVM 模型...")
    svm = SVC(kernel='rbf', C=10)  # 使用 RBF 核函数
    svm.fit(X_train, y_train)

    # 测试 SVM 模型
    print("测试 SVM 模型...")
    y_pred_svm = svm.predict(X_test)
    accuracy_svm = accuracy_score(y_test, y_pred_svm)
    precision_svm = precision_score(y_test, y_pred_svm)
    recall_svm = recall_score(y_test, y_pred_svm)
    f1_svm = f1_score(y_test, y_pred_svm)

    # 列表展示结果
    results = {
        "Model": ["KNN", "SVM"],
        "Accuracy": [accuracy_knn, accuracy_svm],
        "Precision": [precision_knn, precision_svm],
        "Recall": [recall_knn, recall_svm],
        "F1 Score": [f1_knn, f1_svm]
    }

    # 使用 Pandas 输出表格
    results_df = pd.DataFrame(results)
    print("\n分类模型性能评估：")
    print(results_df)

if __name__ == "__main__":
    main()